# import imp
# from time import sleep
# from xml.etree.ElementPath import find
from time import sleep
import cv2
from PIL import Image
# from lightgbm import cv
# from cv2 import circle
import numpy as np
# from torch import square
# from wandb import agent

def show(windowName,img):
    cv2.imshow(windowName,img)
    cv2.waitKey(5)
    cv2.destroyWindow(windowName)
def angle_cos(p0, p1, p2):
    d1, d2 = (p0-p1).astype('float'), (p2-p1).astype('float')
    return abs( np.dot(d1, d2) / np.sqrt( np.dot(d1, d1)*np.dot(d2, d2) ) )
class Points:
    def __init__(self):
        self.bg = cv2.imread("bg.png")
        self.bg_gray  = cv2.cvtColor(self.bg.copy(), cv2.COLOR_BGR2GRAY)
        self.bg_edge = self._edge(self.bg)
        # self._screen(image)
        self.last_agent = []
        self.landmark =self.hough_circle_init()
        self.squares = self.find_squares()
        # if not self.squares:
        #     self.squares = self.find_squares()
        # else:
        #     self.image = cv2.drawContours(self.image,self.squares,-1,(0, 0, 255), 2)
        
    def _screen(self,image):
        # self.image = cv2.imread(image)
        img = Image.fromarray(image[1].astype('uint8')).convert('RGB')
        self.image = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
        cv2.imwrite("test5.png",self.image)
        self.img = self.image.copy()
        self.gray  = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2GRAY)
        self.edge = self._edge(self.image)
        self.image = cv2.drawContours(self.image,self.squares,-1,(0, 0, 255), 2)
        for i in self.landmark:
            cv2.circle(self.image, (i[0], i[1]), i[2], (0, 0, 255), 2)
            cv2.circle(self.image, (i[0], i[1]), 2, (255, 0, 0), 2)
        self.cur_agent = []
        self.agent_circle()
        show("detect",self.image)
    def _edge(self,img):
        blurred = cv2.GaussianBlur(img, (3, 3), 0)  # 高斯模糊
        gray = cv2.cvtColor(blurred,cv2.COLOR_BGR2GRAY)  # 灰路图像
        xgrad = cv2.Sobel(gray, cv2.CV_16SC1, 1, 0)  # xGrodient
        ygrad = cv2.Sobel(gray, cv2.CV_16SC1, 0, 1)  # yGrodient
        edge_output = cv2.Canny(xgrad, ygrad, 100, 150)  # edge
        # show("Canny Edge",edge_output)
        return edge_output
        # #  彩色边缘
        # dst = cv.bitwise_and(image, image, mask=edge_output)
        # cv.imshow("Color Edge", dst)
    # def get_vel(self):
    #     self.cache = []
    #     self.size = 128
    #     self.append(self.agent)
    def agent_circle(self):
        self._agent = cv2.bitwise_xor(self.bg,self.img)
        self._agent = cv2.bitwise_not(self._agent)
        # show("agent",self._agent)
        self.gray_ag = cv2.cvtColor(self._agent, cv2.COLOR_BGR2GRAY)
        circles = cv2.HoughCircles(self.gray_ag, cv2.HOUGH_GRADIENT, 1, 10,  param1=50,param2=30,minRadius=5,maxRadius=25)
        circles = np.uint16(np.around(circles))
        # agent0 = [172,172,236]
        # agent1 = [172,236,172]#[172,172,236]
        for i in circles[0, :]:
            if len(self.last_agent)<2:
                print("保存上一帧信息")
                self.last_agent.append([i[0],i[1],i[2],self.img[i[0],i[1]]])
                # if all(self.image[i[0],i[1]][k]==agent0[k] for k in range(3)):
                #     print("找到红色圆圈")
                #     self.cur_agent[0] = [i[0],i[1],i[2]]
                # elif all(self.image[i[0],i[1]][k]==agent1[k] for k in range(3)):
                #     print("找到绿色圆圈")
                #     self.cur_agent[1] = [i[0],i[1],i[2]]
            self.cur_agent.append([i[0],i[1],i[2],self.img[i[0],i[1]]])
            print("BGR:",self.image[i[0],i[1]])
            cv2.circle(self.image, (i[0], i[1]), i[2], (0, 0, 255), 2)
            cv2.circle(self.image, (i[0], i[1]), 2, (255, 0, 0), 2)
        # print("当前帧",self.cur_agent)
        for i in range(2):
            if len(self.cur_agent)==0:
                self.cur_agent = self.last_agent
            elif len(self.cur_agent)==1:
                if all(self.cur_agent[0][3]==self.last_agent[0][3]):
                    self.last_agent[0] = self.cur_agent[0]
                    self.cur_agent = self.last_agent
                else:
                    self.last_agent[1] = self.cur_agent[0]
                    self.cur_agent = self.last_agent
            else:
                self.last_agent[i] = self.cur_agent[i]
                # print("更新上一帧",self.last_agent)
        for i in range(2):
            cv2.circle(self.image, (self.cur_agent[i][0], self.cur_agent[i][1]), self.cur_agent[i][2], (0, 0, 255), 2)
            cv2.circle(self.image, (self.cur_agent[i][0], self.cur_agent[i][1]), 2, (255, 0, 0), 2)
        # print("agent:",circles[0])
    def hough_circle_init(self):
        # 霍夫圆检测对噪声敏感，边缘检测消噪
        # dst = cv2.pyrMeanShiftFiltering(self.image.copy(), 10, 100)  # 边缘保留滤波EPF
        # show("epf",dst)
        # edge = self._edge(dst)
        # gray = cv2.cvtColor(dst, cv2.COLOR_BGR2GRAY)
        circles = cv2.HoughCircles(self.bg_gray, cv2.HOUGH_GRADIENT, 1, 10,  param1=50,param2=25,minRadius=5,maxRadius=25)
        circles = np.uint16(np.around(circles))  #把circles包含的圆心和半径的值变成整数
        img = self.bg.copy()
        # img2 = self.image.copy()
        check = [204,153,204]#[204,153,204]
        landmark = [64,64,64]
        # agent0 = [172,172,236]
        # agent1 = [172,236,172]#[172,172,236]
        for i in circles[0, :]:
            if all(img[i[0],i[1]][k]==check[k] for k in range(3)):
                print("误判目标",[i[0],i[1]])
                continue
            elif all(img[i[0],i[1]][k]==landmark[k] for k in range(3)):
                print("找到黑色圆圈")
                # self.landmark.append([i[0],i[1]])
            # elif all(img[i[0],i[1]][k]==agent0[k] for k in range(3)):
            #     print("找到红色圆圈")
            #     cv2.circle(img2, (i[0], i[1]), i[2], (0, 0, 0), -1)
            #     self.agent[0] = [i[0],i[1]]
            # elif all(img[i[0],i[1]][k]==agent1[k] for k in range(3)):
            #     print("找到绿色圆圈")
            #     cv2.circle(img2, (i[0], i[1]), i[2], (0, 0, 0), -1)
            #     self.agent[1] = [i[0],i[1]]
            # print("BGR:",img[i[0],i[1]])
            # if i[0]==400:
            #     cv2.circle(img2, (i[0], i[1]), i[2]+2, (255, 255, 255), -1)
            # elif i[0]==600:
                # cv2.circle(img2, (i[0], i[1]), i[2]+2, (255, 255, 255), -1)
            # cv2.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
            # cv2.circle(img, (i[0], i[1]), 2, (255, 0, 0), 2)
        # show("circle image", img)
        # self.bg = img
        # self.bg = img2
        # show("bg image", img2)
        # cv2.imwrite("bg.png",img2)
        print(circles[0,:])
        return circles[0]
    def find_squares(self):
        squares = []
        # img = cv2.GaussianBlur(img, (11,11), 0)   
        # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # bin = cv2.Canny(gray, 0, 255, apertureSize=3) 
        ret,bin = cv2.threshold(self.bg_edge, 127, 255, cv2.THRESH_BINARY)   
        contours, _hierarchy = cv2.findContours(bin, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        # contours, _hierarchy = cv2.findContours(bin, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        # print("轮廓数量：%d" % len(contours))
        # 轮廓遍历
        gpoint = []
        for cnt in contours:
            cnt_len = cv2.arcLength(cnt, True) #计算轮廓周长
            cnt = cv2.approxPolyDP(cnt, 0.02*cnt_len, True) #多边形逼近
            # 条件判断逼近边的数量是否为4，轮廓面积是否大于1000，检测轮廓是否为凸的
            if len(cnt) == 4 and cv2.isContourConvex(cnt):
            # if len(cnt) == 4 and cv2.contourArea(cnt) > 1000 and cv2.isContourConvex(cnt):
                M = cv2.moments(cnt) #计算轮廓的矩
                cx = int(M['m10']/M['m00'])
                cy = int(M['m01']/M['m00'])#轮廓重心
                
                cnt = cnt.reshape(-1, 2)
                max_cos = np.max([angle_cos( cnt[i], cnt[(i+1) % 4], cnt[(i+2) % 4] ) for i in range(4)])
                # 只检测矩形（cos90° = 0）
                if max_cos < 0.1 and [cx,cy] not in gpoint:
                # 检测四边形（不限定角度范围）
                #if True:
                    # index = index + 1
                    # cv2.putText(img,("#%d"%index),(cx,cy),font,0.7,(255,0,255),2)
                    squares.append(cnt)
                    gpoint.append([cx,cy])
        # img = self.image.copy()
        # self.image = cv2.drawContours(img,squares,-1,(0, 0, 255), 2)
        # show("squares",img)
        # print("矩形边界点",squares)
        # print("矩形重心：",gpoint)
        return squares

if __name__=="__main__":
    p = Points()
    p._screen("test.jpg")
    # p.agent_circle()
    # p.get_points()
    # p.edge()
    # p.get_agent()
    # cv2.imshow("test",p.image)
    # cv2.waitKey()
